from typing import Dict, List, Optional, Set
from loguru import logger
from nicegui import ui, app
from tasks.model.task import TaskStatus, TaskSpec
from tasks.server_helper import ServerHelper


def get_status_color(status: TaskStatus) -> str:
    """Get color hex code based on task status."""
    status_value = status.value if isinstance(status, TaskStatus) else status
    return {
        TaskStatus.PENDING.value: "#6B7280",  # gray-500
        TaskStatus.READY.value: "#3B82F6",    # blue-500 (darker)
        TaskStatus.RUNNING.value: "#3B82F6",  # blue-500 (darker)
        TaskStatus.COMPLETED.value: "#10B981", # green-500 (darker)
        TaskStatus.FAILED.value: "#EF4444",   # red-500 (darker)
        TaskStatus.CANCELLED.value: "#6B7280", # gray-500
    }.get(status_value, "#6B7280")


def build_task_tree(
    task_spec: TaskSpec,
    server_helper: ServerHelper,
    visited: Optional[Dict[str, Dict]] = None,
    in_progress: Optional[Set[str]] = None,
) -> Dict:
    """Build a tree node for a task and its dependencies.

    Args:
        task_spec: The TaskSpec to build the tree for
        server_helper: ServerHelper instance to fetch dependent tasks
        visited: Dict of already visited task hashes to their tree nodes
        in_progress: Set of task hashes currently being processed to detect cycles

    Returns:
        Dict containing the tree node data
    """
    if visited is None:
        visited = {}
    if in_progress is None:
        in_progress = set()

    # If we've already built this node, return it
    if task_spec.task_hash in visited:
        return visited[task_spec.task_hash]

    # If we're currently processing this node, we've found a cycle
    if task_spec.task_hash in in_progress:
        logger.warning(f"Circular dependency detected for task {task_spec.task_hash}")
        return None

    in_progress.add(task_spec.task_hash)

    # Get dependent tasks
    dependent_hashes = task_spec.get_dependent_task_hashes()
    children = []

    for dep_hash in dependent_hashes:
        try:
            dep_task = server_helper.get_task_spec(dep_hash)
            child_node = build_task_tree(dep_task, server_helper, visited, in_progress)
            if child_node:
                children.append(child_node)
        except Exception as e:
            logger.warning(f"Failed to get dependent task {dep_hash}: {e}")

    in_progress.remove(task_spec.task_hash)

    # Count unique completed tasks in the entire tree
    def count_unique_tasks(
        nodes: List[Dict], counted: Optional[Set[str]] = None
    ) -> tuple[int, int]:
        if counted is None:
            counted = set()

        completed = 0
        total = 0

        for node in nodes:
            if node["id"] not in counted:
                counted.add(node["id"])
                total += 1
                if node["status"] == TaskStatus.COMPLETED:
                    completed += 1
                child_completed, child_total = count_unique_tasks(
                    node["children"], counted
                )
                completed += child_completed
                total += child_total

        return completed, total

    completed_count, total_count = count_unique_tasks(children)

    # Build node label with task field first, then hash
    label = f"{task_spec.task} ({task_spec.task_hash})"
    if task_spec.name:
        label = f"{task_spec.name} - {label}"
    if total_count > 0:
        label += f" [{completed_count}/{total_count}]"

    # Add HTML color styling
    color = get_status_color(task_spec.status)
    label = f'<span style="color: {color}">{label}</span>'

    node = {
        "id": task_spec.task_hash,
        "label": label,
        "status": task_spec.status,
        "children": children,
    }

    # Store the node in visited dict
    visited[task_spec.task_hash] = node
    return node


def ui_task_tree():
    # Get server_helper instance from the app context
    server_helper: ServerHelper = app.server_helper
    selected = {
        'nodes': list(),
        'repr': '',
    }

    @ui.refreshable
    def ui_task_tree_():
        # Get named tasks
        named_tasks = server_helper.list_named_task_specs(offset=0, limit=50)

        # Build tree data
        tree_data = []

        def on_tick(e):
            selected['nodes'].clear()
            selected['nodes'].extend(e.value)
            selected['repr'] = ', '.join(e.value)
            ui.notify(f"Selected: {e.value}")

        for task in named_tasks:
            node = build_task_tree(task, server_helper, {})
            if node:
                tree_data.append(node)

        # Create tree view
        with ui.tree(tree_data, tick_strategy='strict', on_tick=on_tick).classes("w-full") as tree:

            def on_node_click(e):
                task_id = e.args["id"]
                ui.navigate.to(f"/task_detail/{task_id}")

            tree.on("nodeClick", on_node_click)

        tree.add_slot(
            "default-header",
            """
                <span :props="props" v-html="props.node.label"></span>
            """,
        )
        tree.expand()
        tree.on_select(lambda e: ui.navigate.to(f'/task_detail/{e.value}'))
        

    # Add refresh button
    with ui.row():
        ui.button("Refresh", on_click=ui_task_tree_.refresh)
        ui.checkbox(
            "Auto refresh",
            on_change=lambda e: ui.timer(5, ui_task_tree_.refresh, active=e.value),
        )
    ui.label().bind_text_from(selected, 'repr').classes('w-full')

    ui_task_tree_()


page_config = {
    "nav_entry": [
        {
            "title": "TASK TREE",
            "url": "/task_tree",
            "function": ui_task_tree,
        }
    ]
}
